import os
import random

import chardet
import pandas as pd
from moviepy.audio.io.AudioFileClip import AudioFileClip
from moviepy.video.io.VideoFileClip import VideoFileClip

import pyJianYingDraft as draft
import librosa
import soundfile as sf

from combinations.config import OUTPUT_PATH, DRAFT_PATH
from combinations.conf import logger
from pyJianYingDraft import Font_type, trange, Text_style, Clip_settings, Export_framerate, Export_resolution

# 配置路径常量
AUDIO_SOURCE_DIR = r"D:\模板6_川普锐评\口播音频"
HOT_WORDS_DIR = r"D:\模板6_川普锐评\hot_words"
SRT_SOURCE_DIR = r"D:\模板6_川普锐评\SRT"
IMAGE_FILE_EXTENSIONS = ('.png', '.jpg', '.jpeg')
ORIGINAL_IMAGE_NAMES = [
    "M9V2XHSNWBVTF.jpg", "M9V32H5IRSBOX.jpg",
    "M9V2XVHQ7G82G.jpg", "M9V31K4PENPIG.jpg",
    "M9V2O1TKJ1OSC.jpg"
]


def load_transcript_data() -> pd.DataFrame:
    """
    加载川普评论内容的CSV文件，并返回DataFrame对象

    返回:
        pd.DataFrame: 包含 Contents、Key_words、Number、Titles 四列的数据
    """
    csv_path = r"D:\模板6_川普锐评\文本\Trump0429.csv"

    # 自动检测编码格式
    with open(csv_path, 'rb') as f:
        result = chardet.detect(f.read())

    data_df = pd.read_csv(csv_path, encoding=result['encoding'])
    logger.info(f"成功加载CSV文件，共 {len(data_df)} 行")
    return data_df


def process_audio_file(
        number: str,
        audio_dir: str,
        speed_factor: float = 1
) -> str:
    """
    处理指定编号的音频文件（加速并保存），返回处理后的音频路径

    参数:
        number (str): 音频编号
        audio_dir (str): 音频源目录
        speed_factor (float): 音频加速倍数，默认不加速

    返回:
        str: 处理后的音频路径
    """
    # 获取匹配编号的音频文件
    audio_files = [f for f in os.listdir(audio_dir) if f.split('_')[0] == str(number)]
    if not audio_files:
        logger.warning(f"未找到编号为 {number} 的音频文件")
        return None

    original_audio_path = os.path.join(audio_dir, audio_files[0])
    logger.info(f"匹配到的音频文件: {audio_files[0]}")

    if speed_factor != 1:
        speed_audio_dir = os.path.join(audio_dir, '音频加速')
        new_audio_filename = f"{number}_X{speed_factor}_{os.path.basename(audio_files[0])}"
        new_audio_path = os.path.join(speed_audio_dir, new_audio_filename)

        if not os.path.exists(new_audio_path):
            os.makedirs(speed_audio_dir, exist_ok=True)
            y, sr = librosa.load(original_audio_path, sr=44100)
            y_fast = librosa.effects.time_stretch(y, rate=speed_factor)
            sf.write(new_audio_path, y_fast, sr, format='WAV', subtype='PCM_24')
            logger.info(f"加速后的音频文件已保存到: {new_audio_path}")
        else:
            logger.info(f"加速音频已存在，跳过处理: {new_audio_path}")
        return new_audio_path
    else:
        return original_audio_path


def replace_text_content(script, title: str):
    """
    替换脚本中的文本轨道内容

    参数:
        script: Draft_folder 实例
        title (str): 新的文本内容
    """
    text_track = script.get_imported_track(draft.Track_type.text, index=0)
    script.replace_text(text_track, 0, text=title)
    logger.info(f"替换【文本】title: {title}")


def replace_image_materials(script, keyword_folder: str):
    """
    替换热点词对应的图片素材，最多替换前5张图。

    参数:
        script: draft.Draft_folder 实例
        keyword_folder (str): 热点词对应的图片目录
    """
    if not os.path.exists(keyword_folder):
        logger.warning(f"热点词素材目录不存在: {keyword_folder}")
        return

    image_files = [f for f in os.listdir(keyword_folder) if f.endswith(IMAGE_FILE_EXTENSIONS)]

    def sort_key(filename):
        if filename[0].isdigit():
            return -int(''.join([char for char in filename if char.isdigit()]))
        return random.random()

    image_files.sort(key=sort_key)

    for idx, img_file in enumerate(image_files[:5]):
        img_path = os.path.join(keyword_folder, img_file)
        material_name = ORIGINAL_IMAGE_NAMES[idx]
        script.replace_material_by_name(material_name=material_name, material=draft.Video_material(img_path))
        logger.info(f"已替换图像素材: {material_name} -> {img_path}")


def import_subtitles(script, number: str):
    """
    导入 SRT 字幕文件并应用样式

    参数:
        script: Draft_folder 实例
        number (str): 当前编号
    """
    srt_file_list = [os.path.join(SRT_SOURCE_DIR, f) for f in os.listdir(SRT_SOURCE_DIR)
                     if f.startswith(f"{number}_") and f.endswith(".srt")]

    if not srt_file_list:
        logger.warning(f"未找到编号为 {number} 的 SRT 文件")
        return

    srt_path = srt_file_list[0]
    logger.info(f"找到对应的 SRT 文件: {srt_path}")

    try:
        text_segment = draft.Text_segment("文字", trange("0s", "10s"),
                                          font=Font_type.李李体,
                                          style=Text_style(size=10, color=(0.7, 0.7, 1.0), underline=False, align=1),
                                          clip_settings=Clip_settings(transform_y=-0.8))
        text_segment.add_effect("7127656352410881316")  # 添加花字效果

        script.import_srt(srt_path, track_name="字幕2", style_reference=text_segment,
                          clip_settings=None, font_size=10)
    except Exception as e:
        logger.error(f"SRT 导入失败: {e}")


def export_video(project_name: str, output_dir: str, resolution, framerate):
    """
    导出视频文件

    参数:
        project_name (str): 工程名称
        output_dir (str): 输出目录
        resolution: 分辨率设置
        framerate: 帧率设置
    """
    ctrl = draft.Jianying_controller()
    os.makedirs(output_dir, exist_ok=True)
    output_path = os.path.join(output_dir, f"{project_name}.mp4")
    ctrl.export_draft(project_name, output_path, resolution=resolution,
                      framerate=framerate, zi_mu=False, kou_tu=False)
    logger.info(f"视频导出完成: {output_path}")
    return output_path


def trim_video_based_on_audio(video_path: str, audio_path: str, output_dir: str):
    """
    根据音频时长裁剪视频

    参数:
        video_path (str): 视频文件路径
        audio_path (str): 音频文件路径
        output_dir (str): 输出目录
    """
    audio_clip = AudioFileClip(audio_path)
    adjusted_audio_duration = audio_clip.duration + 1
    audio_clip.close()

    video_clip = VideoFileClip(video_path)
    if video_clip.duration > adjusted_audio_duration:
        trimmed_video = video_clip.subclipped(0, adjusted_audio_duration)
        new_output_path = os.path.join(output_dir, "已剪辑", f"剪裁_{os.path.basename(video_path)}")
        os.makedirs(os.path.dirname(new_output_path), exist_ok=True)
        trimmed_video.write_videofile(new_output_path, codec="libx264")
        logger.info(f"视频剪辑成功，保存路径为：{new_output_path}")
    video_clip.close()


def trump_template_6():
    """
    主函数：批量生成川普风格短视频
    """
    draft_folder = draft.Draft_folder(DRAFT_PATH)
    data_df = load_transcript_data()

    row_index = 2
    while row_index < len(data_df):
        row_data = data_df.iloc[row_index]
        number = row_data['Number']
        title = row_data['Titles']

        logger.info(f"开始处理第 {row_index} 行，编号: {number}")

        # 构建项目名称
        project_name = f"川普锐评_x_{row_index}"

        # 获取关键词素材目录
        keyword_folder = os.path.join(HOT_WORDS_DIR, str(row_data['Key_words']))

        # 获取并处理音频文件
        processed_audio_path = process_audio_file(number, AUDIO_SOURCE_DIR, speed_factor=1)
        if not processed_audio_path:
            logger.warning(f"编号 {number} 音频处理失败，跳过该条目")
            row_index += 1
            continue

        # 创建工程副本
        script = draft_folder.duplicate_as_template("模板6_川普锐评", project_name,
                                                    include_matting=True,
                                                    allow_replace=True)

        # 替换文本内容
        replace_text_content(script, title)

        # 替换图片素材
        replace_image_materials(script, keyword_folder)

        # 替换音频素材
        audio_track = script.get_imported_track(draft.Track_type.audio, index=0)
        script.replace_material_by_seg(track=audio_track, segment_index=0,
                                       material=draft.Audio_material(processed_audio_path),
                                       handle_extend=[draft.Extend_mode.extend_tail])
        logger.info(f"替换【音频】素材: 口播音频.wav -> {processed_audio_path}")

        # 导入字幕
        import_subtitles(script, number)

        # 保存草稿
        script.save()

        # 导出视频
        export_dir = os.path.join(OUTPUT_PATH, "模版6-川普锐评")
        output_video_path = export_video(project_name, export_dir,
                                         resolution=Export_resolution.RES_1080P,
                                         framerate=Export_framerate.FR_24)

        # 裁剪视频
        trim_video_based_on_audio(output_video_path, processed_audio_path, export_dir)

        row_index += 1


if __name__ == "__main__":
    trump_template_6()
